package server

import (
	"context"
	"log"
	"net"
	"sync"

	"github.com/general252/gout/ucolor"
	"github.com/pion/rtcp"
	"github.com/pion/rtp"
)

type RtspServer struct {
	cfg RtspConfig

	rtspListen *net.TCPListener // rtsp listener
	rtpListen  *net.UDPConn     // rtp  listener
	rtcpListen *net.UDPConn     // rtcp listener

	sessionList    map[string]*RtspSession // session
	sessionListMux sync.Mutex

	dispatchHandler *dispatchEventHandler // 触发handles

	ctx        context.Context
	cancelFunc context.CancelFunc
	wg         sync.WaitGroup
}

func NewServer(cfg RtspConfig) (c *RtspServer, err error) {
	c = &RtspServer{
		cfg:         cfg,
		sessionList: map[string]*RtspSession{},
	}
	c.dispatchHandler = newDispatchEventHandler()

	defer func() {
		if err != nil {
			if c.rtspListen != nil {
				_ = c.rtspListen.Close()
			}

			if c.rtpListen != nil {
				_ = c.rtpListen.Close()
			}

			if c.rtcpListen != nil {
				_ = c.rtcpListen.Close()
			}
		}
	}()

	for i := 0; i < 1; i++ {
		// rtsp
		if c.rtspListen, err = net.ListenTCP("tcp4", &net.TCPAddr{
			IP:   net.IPv4zero,
			Port: c.cfg.RtspPort,
		}); err != nil {
			return nil, err
		}

		// rtp
		if c.rtpListen, err = net.ListenUDP("udp4", &net.UDPAddr{
			IP:   net.IPv4zero,
			Port: c.cfg.RtpPort,
		}); err != nil {
			return nil, err
		}

		// rtcp
		if c.rtcpListen, err = net.ListenUDP("udp4", &net.UDPAddr{
			IP:   net.IPv4zero,
			Port: c.cfg.RtcpPort,
		}); err != nil {
			return nil, err
		}
	}

	c.wg.Add(3)

	go c.rtspHandle()
	go c.rtpHandle()
	go c.rtcpHandle()

	c.AddHandle(GetUID(),
		&UnimplementedEventHandler{
			SessionOpen: func(session *RtspSession) {
				c.lockConnection(func() {
					c.sessionList[session.GetUid()] = session
				})
			},
			SessionClose: func(session *RtspSession) {
				c.lockConnection(func() {
					// 删除
					delete(c.sessionList, session.GetUid())
				})
			},
		},
	)

	return c, nil
}

func (c *RtspServer) Close() {
	if c.rtspListen != nil {
		_ = c.rtspListen.Close()
		c.rtspListen = nil
	}

	if c.rtpListen != nil {
		_ = c.rtpListen.Close()
		c.rtpListen = nil
	}

	if c.rtcpListen != nil {
		_ = c.rtcpListen.Close()
		c.rtcpListen = nil
	}

	c.wg.Wait()
}

func (c *RtspServer) AddHandle(id string, h EventHandler) {
	c.dispatchHandler.AddHandle(id, h)
}

func (c *RtspServer) RemoveHandle(id string) {
	c.dispatchHandler.RemoveHandle(id)
}

// getDispatchHandler 事件触发器
func (c *RtspServer) getDispatchHandler() EventHandler {
	return c.dispatchHandler
}

// rtspHandle rtsp routine
func (c *RtspServer) rtspHandle() {
	defer func() {
		c.wg.Done()
	}()

	for {
		tcpConn, err := c.rtspListen.AcceptTCP()
		if err != nil {
			break
		}

		_ = tcpConn.SetReadBuffer(networkBufferSize)
		_ = tcpConn.SetWriteBuffer(networkBufferSize)

		_ = NewConnection(tcpConn, c)
	}
}

// rtpHandle rtp routine
func (c *RtspServer) rtpHandle() {
	var (
		ssrcTrack    = map[uint32]*PusherTrack{}
		cliAddrTrack = map[string]*PusherTrack{}
		mux          sync.RWMutex
	)

	var findTracker = func(ssrc uint32, addr *net.UDPAddr) (*PusherTrack, bool) {
		mux.RLock()
		track, ok := ssrcTrack[ssrc] // find ssrc
		mux.RUnlock()
		if ok {
			return track, ok
		}

		addrString := addr.String()
		mux.RLock()
		track, ok = cliAddrTrack[addrString] // addr find
		mux.RUnlock()
		if ok {
			return track, ok
		}

		// search
		track, ok = c.getRtpPusherTracker(ssrc, addr)
		if ok {
			mux.Lock()
			// save
			if track.ssrc == ssrc {
				ssrcTrack[track.ssrc] = track
			}
			if rtpAddr, yes := track.GetRtpClientAddr(); yes && rtpAddr == addrString {
				cliAddrTrack[addrString] = track
			}
			mux.Unlock()
		}

		return track, ok
	}

	handleUid := GetUID()
	c.AddHandle(handleUid, &UnimplementedEventHandler{
		SessionClose: func(session *RtspSession) {
			session.GetStreamPusher().RangeTrack(func(track *PusherTrack) {
				mux.Lock()
				defer mux.Unlock()

				delete(ssrcTrack, track.ssrc)
				if rtpAddr, yes := track.GetRtpClientAddr(); yes {
					delete(cliAddrTrack, rtpAddr)
				}
			})
		},
	})

	defer func() {
		c.RemoveHandle(handleUid)
		c.wg.Done()
	}()

	var (
		buffer    = make([]byte, 1500)
		udpListen = c.rtpListen
	)

	_ = udpListen.SetWriteBuffer(networkBufferSize)
	_ = udpListen.SetReadBuffer(networkBufferSize)

	for {
		// 读取rtp数据
		n, addr, err := udpListen.ReadFromUDP(buffer)
		if err != nil {
			break
		}

		rtpPacketData := buffer[:n]

		var packet rtp.Packet
		var ssrc uint32
		// 获取ssrc
		{
			if err = packet.Unmarshal(rtpPacketData); err != nil {
				log.Println(err)
				continue
			}
			ssrc = packet.SSRC
		}

		// 根据ssrc/addr查找对应的track
		track, ok := findTracker(ssrc, addr)
		if !ok {
			log.Printf("not found tracker. ssrc: %v address: %v", ucolor.HiYellow(ssrc), ucolor.HiYellow(addr))
			continue
		}

		err = track.WriteRTP(&packet, rtpPacketData)
		if err != nil {
			log.Printf("write RTP fail. %v %v", track.pusher.session.baseURL, err)
		}
	}
}

// rtcpHandle rtcp routine
func (c *RtspServer) rtcpHandle() {
	defer func() {
		c.wg.Done()
	}()

	udpListen := c.rtcpListen
	_ = udpListen.SetWriteBuffer(networkBufferSize)
	_ = udpListen.SetReadBuffer(networkBufferSize)

	buffer := make([]byte, 2048)
	for {
		n, addr, err := udpListen.ReadFromUDP(buffer)
		if err != nil {
			break
		}

		rtcpPacketData := buffer[:n]

		rtcpPackets, err := rtcp.Unmarshal(rtcpPacketData)
		if err != nil {
			log.Println(err)
			continue
		}

		// rtcp, 发送者、接收者都可以发送
		// 接收者报告是player发送的,
		// 发送者报告是pusher发送的

		// player的track最好有个单独的ssrc

		for _, packet := range rtcpPackets {
			switch rtcpPacket := packet.(type) {
			case *rtcp.SenderReport:
				t, ok := c.getRtcpPusherTracker(rtcpPacket.SSRC, addr)
				if ok {
					// 客户端发送接收者报告时, 带过来的ssrc
					t.SetSSRC(rtcpPacket.SSRC)
					t.OnRtcpPacket(&RtcpPacket{
						Packet:         rtcpPacket,
						FromConnection: udpListen,
						FromAddr:       addr,
					})
				} else {
					log.Printf("%v", ucolor.HiRed("not found pusher tracker"))
				}
				log.Printf("sender report. %v(0x%X) from: %v info: %v",
					rtcpPacket.SSRC, rtcpPacket.SSRC, addr.String(), rtcpPacket.String())

			case *rtcp.ReceiverReport:
				t, ok := c.getRtcpPlayerTracker(rtcpPacket.SSRC, addr.String())
				if ok {
					t.OnRtcpPacket(&RtcpPacket{
						Packet:         rtcpPacket,
						FromConnection: udpListen,
						FromAddr:       addr,
					})
				} else {
					log.Printf("%v", ucolor.HiRed("not found player tracker"))
				}
				log.Printf("receiver report. %v 0x%X from: %v, info: %v",
					rtcpPacket.SSRC, rtcpPacket.SSRC, addr.String(), rtcpPacket.String())

				// TypeTransportSpecificFeedback
			case *rtcp.TransportLayerNack:
				log.Println("TypeTransportSpecificFeedback TransportLayerNack")
				t, ok := c.getRtcpPlayerTracker(rtcpPacket.SenderSSRC, addr.String())
				if ok {
					t.OnRtcpPacket(&RtcpPacket{
						Packet:         rtcpPacket,
						FromConnection: udpListen,
						FromAddr:       addr,
					})
				}
			case *rtcp.RapidResynchronizationRequest:
				log.Println("TypeTransportSpecificFeedback RapidResynchronizationRequest")
			case *rtcp.TransportLayerCC:
				log.Println("TypeTransportSpecificFeedback TransportLayerCC")

				// TypePayloadSpecificFeedback
			case *rtcp.PictureLossIndication:
				log.Println("TypePayloadSpecificFeedback PictureLossIndication")
				t, ok := c.getRtcpPlayerTracker(rtcpPacket.SenderSSRC, addr.String())
				if ok {
					t.OnRtcpPacket(&RtcpPacket{
						Packet:         rtcpPacket,
						FromConnection: udpListen,
						FromAddr:       addr,
					})
				}
			case *rtcp.SliceLossIndication:
				log.Println("TypePayloadSpecificFeedback SliceLossIndication")
				t, ok := c.getRtcpPlayerTracker(rtcpPacket.SenderSSRC, addr.String())
				if ok {
					t.OnRtcpPacket(&RtcpPacket{
						Packet:         rtcpPacket,
						FromConnection: udpListen,
						FromAddr:       addr,
					})
				}
			case *rtcp.ReceiverEstimatedMaximumBitrate:
				log.Println("TypePayloadSpecificFeedback ReceiverEstimatedMaximumBitrate")
			case *rtcp.FullIntraRequest:
				log.Println("TypePayloadSpecificFeedback FullIntraRequest")
				// TODO:
			case *rtcp.ExtendedReport:
				log.Println("TypePayloadSpecificFeedback ExtendedReport")
			case *rtcp.RawPacket:
				log.Println("RawPacket")
			}
		}
	}
}

func (c *RtspServer) GetRtpPort() int {
	return c.cfg.RtpPort
}

func (c *RtspServer) GetRtcpPort() int {
	return c.cfg.RtcpPort
}

func (c *RtspServer) GetRtpSender() *net.UDPConn {
	return c.rtpListen
}

func (c *RtspServer) GetRtcpSender() *net.UDPConn {
	return c.rtcpListen
}

func (c *RtspServer) GetPusher(urlPath string) (*StreamPusher, bool) {
	var r *StreamPusher

	c.lockConnection(func() {
		for _, conn := range c.sessionList {
			if conn.GetConnectionType() != ConnectionTypePusher {
				continue
			}

			object := conn.GetStreamPusher()
			if object.GetPath() == urlPath {
				r = object
				break
			}
		}
	})

	return r, r != nil
}

// getRtpPusherTracker 根据ssrc/addr查找pusher tracker
func (c *RtspServer) getRtpPusherTracker(ssrc uint32, addr *net.UDPAddr) (*PusherTrack, bool) {
	c.sessionListMux.Lock()
	defer c.sessionListMux.Unlock()

	for _, session := range c.sessionList {
		if !session.IsPusher() {
			continue
		}

		for _, pusherTrack := range session.mStreamPusher.Tracks {
			if pusherTrack.ssrc == ssrc {
				return pusherTrack, true
			}
		}
	}

	addrString := addr.String()
	for _, session := range c.sessionList {
		if !session.IsPusher() {
			continue
		}

		for _, pusherTrack := range session.mStreamPusher.Tracks {
			if rtpAddr, _ := pusherTrack.GetRtpClientAddr(); rtpAddr == addrString {
				return pusherTrack, true
			}
		}
	}

	return nil, false
}

func (c *RtspServer) getRtcpPusherTracker(ssrc uint32, addr *net.UDPAddr) (*PusherTrack, bool) {
	c.sessionListMux.Lock()
	defer c.sessionListMux.Unlock()

	for _, session := range c.sessionList {
		if !session.IsPusher() {
			continue
		}

		for _, pusherTrack := range session.mStreamPusher.Tracks {
			if pusherTrack.ssrc == ssrc {
				return pusherTrack, true
			}
		}
	}

	addrString := addr.String()
	for _, session := range c.sessionList {
		if !session.IsPusher() {
			continue
		}

		for _, pusherTrack := range session.mStreamPusher.Tracks {
			if rtcpAddr, ok := pusherTrack.GetRtcpClientAddr(); ok && rtcpAddr == addrString {
				return pusherTrack, true
			}
		}
	}

	return nil, false
}

func (c *RtspServer) getRtcpPlayerTracker(ssrc uint32, addr string) (*PlayerTrack, bool) {
	c.sessionListMux.Lock()
	defer c.sessionListMux.Unlock()

	for _, session := range c.sessionList {
		if !session.IsPlayer() {
			continue
		}

		for _, playerTrack := range session.mStreamPlayer.Tracks {
			if playerTrack.ssrc == ssrc {
				return playerTrack, true
			}

			if rtcpAddr, _ := playerTrack.GetRtcpClientAddr(); rtcpAddr == addr {
				return playerTrack, true
			}
		}
	}

	return nil, false
}

func (c *RtspServer) lockConnection(f func()) {
	c.sessionListMux.Lock()
	defer c.sessionListMux.Unlock()

	f()
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

type RtspConfig struct {
	RtspPort int // rtsp tcp port
	RtpPort  int // rtp udp port
	RtcpPort int // rtcp udp port
}
